'''
Created on Feb 16, 2011

@author: oabalbin
This script do not perform snps functional
analysis. Instead, the scripts prepares
files for batch query of polyphen and SIFT
based on an already annotated vcf or log file. 
'''

from optparse import OptionParser
from exome.vcftools.vcf import *
from exome.vcftools.write_vcf import write_vcf_line
from collections import defaultdict


def make_polyphen2_batch_file(input_file_name, vcf_outfile_name):
    vcOutfile=open(vcf_outfile_name,'w')
    thisvcf=VCF()
    thisvcf.readFromFile(input_file_name)
    SNV_dict=defaultdict()
    
    for snp in thisvcf:
        site=snp['COORD'].split('@')
        CHROM, POS=str(site[0]),str(site[1])
        LOC=CHROM+':'+POS
        SNV=",".join(snp['REF']+snp['ALT']).replace(',','/')
        snvinfo = [LOC,SNV]
        SNV_dict[LOC]=snvinfo
    
    SNV_list=SNV_dict.keys()
    SNV_list.sort()
    for snv in SNV_list:
        line=SNV_dict[snv]
        vcOutfile.write(",".join(line).replace(',','\t')+'\n')

    pass


def read_polyphen2_output(file_name):
    '''
    '''
    inputfile=open(file_name)
    pp2_sites=defaultdict()
    
    for l in inputfile:
        if l.startswith('#'):
            continue
        fields=l.strip('\n').split('\t')
        
        o_snp_id=fields[0].split('.')[0]
        POS=o_snp_id.replace(':','@')
        prot_name=fields[2]
        prediction=fields[6]
        pp2_sites[POS]=prediction.strip(' ').replace(' ','_')
    
    return pp2_sites
    


def annotate_vcf_with_polyphen2(input_file_name, vcf_outfile_name, annotated_sites):
    '''
    '''
    print >> sys.stderr,"Annotating snps according to Poly-phen2"
    PP2='PP2'
    vcOutfile=open(vcf_outfile_name,'w')
    thisvcf=VCF()
    thisvcf.readFromFile(input_file_name)
    
    # update info tags.
    thisvcf.infotags['_INFOTAGS'].append(PP2)
    i=0
    
    print >> sys.stderr,"Annotated sites with polyphen %d"%len(annotated_sites)
    for coord, pp2_info in annotated_sites.iteritems():
        coord=[coord]
        try:
            thissnps= thisvcf.get(coord)
            for snp in thissnps:
                snp[PP2]=pp2_info
                line = write_vcf_line(snp, thisvcf.infotags['_INFOTAGS'], thisvcf.infotags['_SAMPLETAGS'])
                vcOutfile.write(line+'\n')    
                i+=1
        except:
            print >> sys.stderr,"This annotated SNV % were not found in the vcf"%str(coord)
            
                
    print >> sys.stderr,"Done. %d SNVs were annotatoted according to polyphen2"%i
    
    
    
input_file_name='/exds/users/oabalbin/projects/snps/exomes/FMPN_80/analysis/isec_calls/FMPN.t.snps.common.pipes.shet.vcf'
vcf_outfile_name=input_file_name.replace('.vcf','.pp2.batch')
#make_polyphen2_batch_file(input_file_name, vcf_outfile_name)

##
polyphen_file='/exds/users/oabalbin/projects/snps/exomes/FMPN_80/analysis/isec_calls/FMPN.t.common.pipes.het.pph2-short.txt'
input_file_name2='/exds/users/oabalbin/projects/snps/exomes/FMPN_80/analysis/isec_calls/FMPN.t.snps.common.pipes.het.annotated.vcf'
vcf_outfile_name2=input_file_name.replace('.vcf','.pp2.annot')

pp2_sites = read_polyphen2_output(polyphen_file)

annotate_vcf_with_polyphen2(input_file_name2,vcf_outfile_name2,pp2_sites)
                            
'''
if __name__ == '__main__':
        
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("--local", dest="local", action="store_true", default=False)
    optionparser.add_option("--cluster", dest="cluster", action="store_true", default=False)
    optionparser.add_option("-f", "--vcf_file", dest="vcf_file",
                            help="vcf_file")
    optionparser.add_option("-o", "--out_vcf_file", dest="out_vcf_file",
                            help="out_vcf_file")
    
    (options, args) = optionparser.parse_args()    
''' 
    

    
